/*
 * (c) 2008-2011 Daniel Halperin <dhalperi@cs.washington.edu>
 */
#include "iwl_connector.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <linux/netlink.h>
#include <errno.h>

#define MAX_PAYLOAD 2048
#define SLOW_MSG_CNT 1
int WRITE_TO_FILE = 0;
int NET_TO_FILE = 0;
int sock_fd = -1; // the socket
int csi_sock_fd = -1;
int csi_fd = -1;

FILE *out = NULL;

void check_usage(int argc, char **argv);

FILE *open_file(char *filename, char *spec);

void caught_signal(int sig);

void exit_program(int code);
void exit_program_err(int code, char *func);

int main(int argc, char **argv)
{
    /* Local variables */
    struct sockaddr_nl proc_addr, kern_addr; // addrs for recv, send, bind

    struct sockaddr_in csi_addr; //csi_addr for sending data to the experiment

    struct cn_msg *cmsg;
    char buf[MAX_PAYLOAD];
    char filename[260];
    int ret;
    unsigned short l, l2;
    int count = 0;
    int count1 = 0;

    /* Make sure usage is correct */
    check_usage(argc, argv);

    /* Open and check log file */
    if (WRITE_TO_FILE)
        out = open_file(argv[1], "w");

    /* Setup the socket */
    sock_fd = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_CONNECTOR);
    csi_sock_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (sock_fd == -1)
        exit_program_err(-1, "socket");
    if (csi_sock_fd == -1)
        exit_program_err(-1, "csi_socket");
    /* Initialize the address structs */
    memset(&proc_addr, 0, sizeof(struct sockaddr_nl));
    proc_addr.nl_family = AF_NETLINK;
    proc_addr.nl_pid = getpid(); // this process' PID
    proc_addr.nl_groups = CN_IDX_IWLAGN;
    memset(&kern_addr, 0, sizeof(struct sockaddr_nl));
    kern_addr.nl_family = AF_NETLINK;
    kern_addr.nl_pid = 0; // kernel
    kern_addr.nl_groups = CN_IDX_IWLAGN;
    memset(&csi_addr, 0, sizeof(struct sockaddr_in));
    csi_addr.sin_family = AF_INET;
    csi_addr.sin_addr.s_addr = 0;
    csi_addr.sin_port = htons(8899);
    int on = 1;
    setsockopt(csi_sock_fd, 1, 2, &on, sizeof(on));
    /* Now bind the socket */
    if (bind(sock_fd, (struct sockaddr *)&proc_addr, sizeof(struct sockaddr_nl)) == -1)
        exit_program_err(-1, "bind");

    /* And subscribe to netlink group */
    {
        int on = proc_addr.nl_groups;
        ret = setsockopt(sock_fd, 270, NETLINK_ADD_MEMBERSHIP, &on, sizeof(on));
        if (ret)
            exit_program_err(-1, "setsockopt");
    }

    /* Set up the "caught_signal" function as this program's sig handler */
    signal(SIGINT, caught_signal);
    signal(SIGPIPE, SIG_IGN);
    struct sockaddr_in addr_in;

    int clientlen;
    if (!WRITE_TO_FILE)
    {
        if (bind(csi_sock_fd, (struct sockaddr *)&csi_addr, sizeof(struct sockaddr)) == -1)
        {
            exit_program_err(-1, "bind_csi");
        }

        listen(csi_sock_fd, 1);
    ACCEPT:
        if (out != NULL)
        {
            fclose(out);
            out = NULL;
        }
        printf("Input the fillename if you want to record the data to a file:");
        scanf("%s", filename);
        if (strlen(filename) != 0)
        {
            printf("Record to file : %s\n",filename);
            out = open_file(filename, "w");
            NET_TO_FILE = 1;
        }
        else
        {
            printf("Ignore recording.\n");
            out = NULL;
            NET_TO_FILE = 0;
        }

        printf("Waiting for client connecting!\n");
        count = 0;
        csi_fd = accept(csi_sock_fd, (struct sockaddr *)&addr_in, (socklen_t *)&clientlen);

        printf("Connected from %s:%d\n", inet_ntoa(addr_in.sin_addr), addr_in.sin_port);
    }

    /* Poll socket forever waiting for a message */
    while (1)
    {

        /* Receive from socket with infinite timeout */
        ret = recv(sock_fd, buf, sizeof(buf), 0);
        if (ret == -1)
            continue;
        /* Pull out the message portion and print some stats */
        cmsg = NLMSG_DATA(buf);
        if (count % SLOW_MSG_CNT == 0)
        {
        }
        //printf("received %d bytes: id: %d val: %d seq: %d clen: %d\n", cmsg->len, cmsg->id.idx, cmsg->id.val, cmsg->seq, cmsg->len);
        /* Log the data to file */
        l = (unsigned short)cmsg->len;
        l2 = htons(l);

        if (WRITE_TO_FILE)
        {
            fwrite(&l2, 1, sizeof(unsigned short), out);
            ret = fwrite(cmsg->data, 1, l, out);

            if (ret != l)
                exit_program_err(1, "fwrite");
        }
        else
        {
            int i;
            unsigned short l2tmp = htons(l + 5);
            unsigned char buffer[MAX_PAYLOAD];
            unsigned char header[] = {0xaa, 0x55};
            unsigned char lrc = 0;
            memcpy(buffer, header, 2);
            memcpy(buffer + 2, &l2tmp, 2);
            memcpy(buffer + 4, cmsg->data, l);
            if (6 + l < MAX_PAYLOAD)
            {
                for (i = 0; i < 4 + l; i++)
                {
                    lrc ^= buffer[i];
                }
                buffer[4 + l] = lrc;
                ret = send(csi_fd, buffer, 5 + l, 0);
                if (ret != 5 + l)
                {
                    //FAILED
                    goto ACCEPT;
                    
                }
            }
        }
        if(NET_TO_FILE){
            fwrite(&l2, 1, sizeof(unsigned short), out);
            ret = fwrite(cmsg->data, 1, l, out);

            if (ret != l)
                exit_program_err(1, "fwrite");
        }
        if (count % 100 == 0)
            printf("message count = %u\n", count);
        ++count;
        usleep(500);
    }
    exit_program(0);
    return 0;
}

void check_usage(int argc, char **argv)
{
    if (argc < 3)
    {
        WRITE_TO_FILE = argc - 1;
    }
    else
    {
        fprintf(stderr, "Usage: %s [<output_file>]\n", argv[0]);
        exit_program(1);
    }
}

FILE *open_file(char *filename, char *spec)
{
    FILE *fp = fopen(filename, spec);
    if (!fp)
    {
        perror("fopen");
        exit_program(1);
    }
    return fp;
}

void caught_signal(int sig)
{
    fprintf(stderr, "Caught signal %d\n", sig);
    exit_program(0);
}

void exit_program(int code)
{
    if (out)
    {
        fclose(out);
        out = NULL;
    }
    if (sock_fd != -1)
    {
        close(sock_fd);
        sock_fd = -1;
    }
    exit(code);
}

void exit_program_err(int code, char *func)
{
    perror(func);
    exit_program(code);
}
